package com.sunday.common.mq.kafka.study.spring.e21_Brave_Trace_Pull;

import brave.Tracing;
import brave.kafka.clients.KafkaTracing;
import brave.messaging.MessagingTracing;
import com.sunday.common.mq.kafka.brave.core.TracingKafkaConsumerFactory;
import com.sunday.common.mq.kafka.brave.core.TracingKafkaProducerFactory;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.kafka.DefaultKafkaConsumerFactoryCustomizer;
import org.springframework.boot.autoconfigure.kafka.DefaultKafkaProducerFactoryCustomizer;
import org.springframework.boot.autoconfigure.kafka.KafkaConnectionDetails;
import org.springframework.boot.autoconfigure.kafka.KafkaProperties;
import org.springframework.boot.context.properties.PropertyMapper;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.kafka.core.KafkaTemplate;
import org.springframework.kafka.core.ProducerFactory;
import org.springframework.kafka.support.ProducerListener;
import org.springframework.kafka.support.converter.RecordMessageConverter;

import java.util.Map;

@Slf4j
@Configuration
@RequiredArgsConstructor
public class KafkaBraveConfiguration {

    private final KafkaProperties properties;

    @Bean
    public MessagingTracing messagingTracing(Tracing tracing) {
        return MessagingTracing.create(tracing);
    }

    @Bean
    public KafkaTracing springRocketTracing(MessagingTracing messagingTracing) {
        return KafkaTracing.create(messagingTracing);
    }

    /**
     * Producer 埋点
     */
    @Bean
    public TracingKafkaProducerFactory<?, ?> kafkaProducerFactory(KafkaConnectionDetails connectionDetails, ObjectProvider<DefaultKafkaProducerFactoryCustomizer> customizers, KafkaTracing kafkaTracing) {
        Map<String, Object> properties = this.properties.buildProducerProperties();
        properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, connectionDetails.getProducerBootstrapServers());
        TracingKafkaProducerFactory<?, ?> factory = new TracingKafkaProducerFactory<>(properties, kafkaTracing);
        String transactionIdPrefix = this.properties.getProducer().getTransactionIdPrefix();
        if (transactionIdPrefix != null) {
            factory.setTransactionIdPrefix(transactionIdPrefix);
        }
        customizers.orderedStream().forEach((customizer) -> customizer.customize(factory));
        return factory;
    }


    /**
     * consumer 埋点
     */
    @Bean
    public TracingKafkaConsumerFactory<?, ?> kafkaConsumerFactory(KafkaConnectionDetails connectionDetails, ObjectProvider<DefaultKafkaConsumerFactoryCustomizer> customizers, KafkaTracing kafkaTracing) {
        Map<String, Object> properties = this.properties.buildConsumerProperties();
        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, connectionDetails.getConsumerBootstrapServers());
        TracingKafkaConsumerFactory<?, ?> factory = new TracingKafkaConsumerFactory<>(properties, kafkaTracing);
        customizers.orderedStream().forEach((customizer) -> customizer.customize(factory));
        /**
         * pull
         */
        factory.updateConfigs(Map.of(
                ConsumerConfig.GROUP_ID_CONFIG, "myId_15_B",
                ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1
        ));
        return factory;
    }

    /**
     * Consumer Pull 埋点
     */
    @Bean
    @ConditionalOnMissingBean(KafkaTemplate.class)
    public KafkaTemplate<?, ?> kafkaTemplate(ProducerFactory<Object, Object> kafkaProducerFactory,
                                             ProducerListener<Object, Object> kafkaProducerListener,
                                             ObjectProvider<RecordMessageConverter> messageConverter,
                                             TracingKafkaConsumerFactory<Object, Object> kafkaConsumerFactory
    ) {
        PropertyMapper map = PropertyMapper.get().alwaysApplyingWhenNonNull();
        KafkaTemplate<Object, Object> kafkaTemplate = new KafkaTemplate<>(kafkaProducerFactory);
        messageConverter.ifUnique(kafkaTemplate::setMessageConverter);
        map.from(kafkaProducerListener).to(kafkaTemplate::setProducerListener);
        map.from(this.properties.getTemplate().getDefaultTopic()).to(kafkaTemplate::setDefaultTopic);
        map.from(this.properties.getTemplate().getTransactionIdPrefix()).to(kafkaTemplate::setTransactionIdPrefix);
        /**
         * pull
         */
        kafkaTemplate.setConsumerFactory(kafkaConsumerFactory);
        return kafkaTemplate;
    }

}
